import os
import numpy as np
from app import conf
if os.name == 'posix' and 'Linux' in os.uname().sysname:
	__import__('pysqlite3')
	import sys
	sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
	
import chromadb

client = chromadb.PersistentClient(os.path.join(conf.DATA_ROOT, 'chromadb'))
collection = client.get_or_create_collection(name='key_collection', metadata={"hnsw:space": "cosine"})


def add(ids, metadatas, embeddings):
	collection.add(
		ids=ids,
		embeddings=embeddings,
		metadatas=metadatas
	)
	
	
def add2(ids, embeddings):
	metadatas = [{'key': id_} for id_ in ids]
	collection.add(
		ids=ids,
		embeddings=embeddings,
		metadatas=metadatas
	)
	
	
def update(ids, metadatas, embeddings):
	collection.update(
		ids=ids,
		embeddings=embeddings,
		metadatas=metadatas
	)
	
	
def delete(ids):
	collection.delete(ids)
	
	
def get_embedding(id_: str):
	return collection.get(id_, include=['embeddings'])['embeddings']
	

def get(id_: str, topK=2):
	results = collection.get(id_, include=['embeddings'])
	results = collection.query(
		query_embeddings=results['embeddings'],
		# query_texts=  # 添加的时候传的是documents，不传embeddings，embeddings由模型自己生成，可用
		n_results=topK,
		where={"key": {"$ne": id_}},  # 不查自己
		# where_document={"$contains": "key-1"}  # optional filter
		include=['embeddings']
	)
	return results


if __name__ == '__main__':
	# keys = [f"key-{i}" for i in range(20)]
	# metadatas = [{'key': k} for k in keys]
	# a = np.random.randn(10, 5)
	# embeddings = np.concatenate([a, a]).tolist()
	
	# add(keys, metadatas, embeddings)
	
	results = get("key-2", 3)
	print("results:", results)
	
	b = 21
	for i in range(5):
		r = get_embedding(f'key-{i}')
		add2([f'key-{b + i}'], r)
	